# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import os.path as osp

import torch
from PIL import Image
from torch.utils.data import Dataset
import torchvision.transforms as T

def build_transforms(cfg):
    normalize_transform = T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD)
    transform = T.Compose([
        T.Resize(cfg.INPUT.SIZE_TEST),
        T.ToTensor(),
        normalize_transform
    ])

    return transform
def read_image(img_path):
    """Keep reading image until succeed.
    This can avoid IOError incurred by heavy IO process."""
    got_img = False
    if not osp.exists(img_path):
        raise IOError("{} does not exist".format(img_path))
    while not got_img:
        try:
            img = Image.open(img_path).convert('RGB')
            got_img = True
        except IOError:
            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
            pass
    return img

import os
class ImageDataset(Dataset):
    """Image Person ReID Dataset"""

    def __init__(self, dir=None, cfg=None):
        transform = build_transforms(cfg)
        self.dataset = os.listdir(dir)
        self.transform = transform
        self.dir = dir

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        # add by fby
        img_path= os.path.join(self.dir, self.dataset[index])
        img = read_image(img_path)

        if self.transform is not None:
            img = self.transform(img)
        img = torch.unsqueeze(img, 0)
        return img, img_path

